import torch
from d2l import torch as d2l
from torch import nn

def normal_weights(m):
        if type(m) == nn.Linear:
            nn.init.normal_(m.weight, std=0.01)
            print("normal",m.weight)

def eye_weights(m):

        if type(m) == nn.Linear:
            nn.init.eye_(m.weight)
            print("eye",m.weight)

def uniform_weights(m):

        if type(m) == nn.Linear:
            nn.init.uniform_(m.weight, a=-1,b=1)
            print("uniform",m.weight)

def train_diffweight(w):
    net = nn.Sequential(nn.Flatten(),nn.Linear(784, 256),nn.ReLU(),nn.Linear(256, 10))
    if w==0:
        net.apply(normal_weights);
    elif w==1:
        net.apply(eye_weights);
    else:
        net.apply(uniform_weights);
    batch_size, lr, num_epochs = 256, 0.1, 10
    loss = nn.CrossEntropyLoss(reduction='none')
    trainer = torch.optim.SGD(net.parameters(), lr=lr)
    train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)
    d2l.train_ch3(net, train_iter, test_iter, loss, num_epochs, trainer)
    d2l.plt.ylabel("Y")


train_diffweight(0)
train_diffweight(1)
train_diffweight(2)